package com.lunatic.assistant.config;


import com.alibaba.cloud.ai.autoconfigure.dashscope.DashScopeChatProperties;
import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeCloudStore;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetriever;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeDocumentRetrieverOptions;
import com.alibaba.cloud.ai.dashscope.rag.DashScopeStoreOptions;
import com.lunatic.assistant.service.impl.MockOrderService;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.rag.retrieval.search.DocumentRetriever;
import org.springframework.ai.reader.TextReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.CommandLineRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Description;
import org.springframework.core.io.Resource;

import java.util.List;
import java.util.function.Function;

@Configuration
public class AiConfiguration {

    @Bean
    public ChatMemory chatMemory(){
        return new InMemoryChatMemory();
    }


//    @Bean
//    public CommandLineRunner commandLineRunner(VectorStore vectorStore,
//                                               @Value("classpath:rag/terms-of-service.txt") Resource termsOfServiceDocs) {
//        return args -> {
//            // Ingest the document into the vector store
//            vectorStore.write(new TokenTextSplitter()
//                    .transform(new TextReader(termsOfServiceDocs).read())
//            );
//        };
//    }

    /**
     * 向量数据库，也可以使用redis等作为向量数据库
     *
     * @param embeddingModel 嵌入模型
     * @return
     */
//    @Bean
//    public VectorStore vectorStore(EmbeddingModel embeddingModel){
//        return SimpleVectorStore.builder(embeddingModel).build();
//    }


    @Bean
    public VectorStore vectorStore(EmbeddingModel embeddingModel) {
        SimpleVectorStore simpleVectorStore = SimpleVectorStore.builder(embeddingModel).build();

        // 生成一个机器人产品说明书的文档
        List<Document> documents = List.of(
                new Document("产品说明书:产品名称：智能机器人\n" +
                        "产品描述：智能机器人是一个智能设备，能够自动完成各种任务。\n" +
                        "功能：\n" +
                        "1. 自动导航：机器人能够自动导航到指定位置。\n" +
                        "2. 自动抓取：机器人能够自动抓取物品。\n" +
                        "3. 自动放置：机器人能够自动放置物品。\n"));

        simpleVectorStore.add(documents);
        return simpleVectorStore;
    }


    @Bean
    public DocumentRetriever documentRetriever(DashScopeApi dashScopeApi) {
        return new DashScopeDocumentRetriever(dashScopeApi,
                DashScopeDocumentRetrieverOptions.builder()
                        .withIndexName("spring-ai知识库")
                        .build());
    }


    @Bean
    public DashScopeCloudStore dashScopeCloudStore(DashScopeApi dashScopeApi) {
        return new DashScopeCloudStore(
                dashScopeApi, new DashScopeStoreOptions("spring-ai知识库"));
    }

}
